#!/usr/bin/env python
import flowd
import string
import getopt
import sys
import MySQLdb
#import psycopg

mask_fields = [
	'TAG'			  ,
	'UNIX_SECS'       ,
    'UNIX_NSECS'      ,
    'SYSUPTIME'       ,
    'EXADDR'          ,   
    'DPKTS'           ,
    'DOCTETS'         ,
    'FIRST'           ,  
    'LAST'            ,
    'ENGINE_TYPE'     ,
    'ENGINE_ID'       ,    
    'SRCADDR'         ,
    'DSTADDR'         ,
    'NEXTHOP'         ,
    'INPUT'           ,
    'OUTPUT'          ,
    'SRCPORT'         ,   
    'DSTPORT'         ,
    'PROT'            ,
    'TOS'             ,
    'TCP_FLAGS'       ,    
    'SRC_MASK'        ,
    'DST_MASK'        ,
    'SRC_AS'          ,
    'DST_AS'          ,   
]

def usage():
	print >> sys.stderr, "export.py (flowd.py version %s)" % \
	    flowd.__version__
	print >> sys.stderr, "Usage: export.py [options] [flowd-store]";
	print >> sys.stderr, "Options:";
	print >> sys.stderr, "      -h       Display this help";
	print >> sys.stderr, "      -f       Format of output. 0 - ASCII, 1 - MySQL, 2 - PostgreSQL. Default 0";
	print >> sys.stderr, "      -m       mask_fields. Select fields from: %s. Default all fields" % ', '.join(mask_fields);
	print >> sys.stderr, "      -o       output file for ASCII format";
	print >> sys.stderr, "      -u       user:password:host:port:name:table Configure MySQL or PostgresSQL access";
	sys.exit(1);

def get_mask_value(mask, flow):
	mask_value = []
	if 'TAG' in mask:
		mask_value.append("%d"%flow.tag)
	if 'UNIX_SECS' in mask:
		mask_value.append("%d"%flow.recv_sec)
	if 'UNIX_NSECS' in mask:
		mask_value.append("%d"%flow.recv_usec)
	if 'SYSUPTIME'  in mask:
		mask_value.append("%d"%flow.sys_uptime_ms)
	if 'EXADDR' in mask: 
		mask_value.append(flow.agent_addr)		
	if 'DPKTS' in mask:
		mask_value.append("%d"%flow.packets)
	if 'DOCTETS' in mask:
		mask_value.append("%d"%flow.octets)
	if 'FIRST' in mask: 
		mask_value.append("%d"%flow.flow_start)
	if 'LAST' in mask:
		mask_value.append("%d"%flow.flow_finish)
	if 'ENGINE_TYPE' in mask:
		mask_value.append("%d"%flow.engine_type)
	if 'ENGINE_ID' in mask:    
		mask_value.append("%d"%flow.engine_id)
	if 'SRCADDR' in mask:
		mask_value.append(flow.src_addr)
	if 'DSTADDR' in mask:
		mask_value.append(flow.dst_addr)
	if 'NEXTHOP' in mask:
		mask_value.append(flow.gateway_addr)
	if 'INPUT' in mask:
		mask_value.append("%d"%flow.if_ndx_in)
	if 'OUTPUT' in mask:
		mask_value.append("%d"%flow.if_ndx_out)
	if 'SRCPORT' in mask:   
		mask_value.append("%d"%flow.src_port)
	if 'DSTPORT' in mask:
		mask_value.append("%d"%flow.dst_port)
	if 'PROT' in mask:
		mask_value.append("%d"%flow.protocol)
	if 'TOS' in mask:
		mask_value.append("%d"%flow.tos)
	if 'TCP_FLAGS' in mask:    
		mask_value.append("%s"%flow.tcp_flags)
	if 'SRC_MASK' in mask:
		mask_value.append("%d"%flow.src_mask)
	if 'DST_MASK'in mask:
		mask_value.append("%d"%flow.dst_mask)
	if 'SRC_AS' in mask:
		mask_value.append("%d"%flow.src_as)
	if 'DST_AS' in mask:  
		mask_value.append("%d"%flow.dst_as)
	
	return mask_value

def create_query(mask, dbtable):
	str = "CREATE TABLE IF NOT EXISTS %s (" %dbtable
	for m in mask:
		if str[-1] != '(':
			str +=", "
		str += m
		if m in ("SRCADDR", "DSTADDR", "EXADDR", "NEXTHOP"):
			str += " VARCHAR(64)"
		else:
			str += " INTEGER"
		
	str += ");"
	return str
	
def format0(outfile, mask, flow):
	outfile.write("%s\n"%",".join(get_mask_value(mask,flow)))

def format1(conn, dbtable, mask, flow):
	cursor = conn.cursor()
	str = "INSERT INTO %s " %dbtable
	str += "(%s) "%','.join(mask)
	str += "VALUES ('%s')"%"','".join(get_mask_value(mask,flow))
	cursor.execute(str)

	
def main():

	try:
		opts, args = getopt.getopt(sys.argv[1:], 'f:m:u:o:h')
	except getopt.GetoptError:
		print >> sys.stderr, "Invalid commandline arguments"
		usage()

	format = None
	conn = None
	mask = None
	outfile = None
	uopt = None
	for o, a in opts:
		if o in ('-h', '--help'):
			usage()
			sys.exit(0)
		if o in ('-f', '--format'):
			format = a
		if o in('-m'):
			mask = a.split(',')
			for m in mask:
				if m not in mask_fields:
					print >> sys.stderr, "Unrecognized field:"+m
					sys.exit(1);
		if o in ('-u'):
			if format not in ('1','2'):
				print >> sys.stderr, "Option -u must be used only for MySQL or PostgresSQL format"
				usage()
			else:
				uopt = 1
				if len(a.split(':')) == 6:
					dbuser, dbpswd, dbhost, dbport, dbname, dbtable = a.split(':')
					dbport = int (dbport)
				else:
					print >> sys.stderr, "Option -u need six values"
					usage()
		if o in ('-o'):
			if format != '0':
				print >> sys.stderr, "Option -o must be used only for ASCII format"
				usage()
			else:
				outfile = open(a, "w")
		
	if len(args) == 0:
		print >> sys.stderr, "No logfile specified"
		usage()
	if format is None:
		format = '0'
	if mask is None:
		mask = mask_fields
	else:
		mask_sort = [m for m in mask_fields if m in mask]
		mask = mask_sort
	if not format.isdigit() or int(format) > 2:
		print >> sys.stderr, "Invalid format"
		usage()
	else:
		if format == '0':
			outfile.write("#:%s\n"%",".join(mask))
		elif format == '1':
			if not uopt:
				print >> sys.stderr, "No -u option specified"
				usage()
			try:
				conn = MySQLdb.Connect(host=dbhost, port=dbport, user=dbuser, passwd=dbpswd, db=dbname)
				cursor = conn.cursor()
				cursor.execute(create_query(mask, dbtable))
			except MySQLdb.Error, e:
				print("Error: unavailable connection with database %s" % dbname)
				sys.exit(-1)	
		elif format == '2':
			if not uopt:
				print >> sys.stderr, "No -u option specified"
				usage()
			try:
				#conn = psycopg.connect('dbname=%s host=%s port=%s user=%s password=%s'%(dbname,dbhost,dbport,dbuser,dbpswd))
				conn = psycopg.connect('dbname=%s user=%s'%(dbname,dbuser,))
				cursor = conn.cursor()
				cursor.execute(string.replace(create_query(mask, dbtable),'IF NOT EXISTS',''))
				#print string.replace(create_query(mask, dbtable),'IF NOT EXISTS','')
			except MySQLdb.Error, e:
				print("Error: unavailable connection with database %s" % dbname)
				sys.exit(-1)
	count = 0
	for ffile in args:
		flog = flowd.FlowLog(ffile, "rb")	
		for flow in flog:
			count += 1
			if format == '0':
				format0(outfile, mask, flow)
			elif format in ('1','2'):
				format1(conn, dbtable, mask, flow)
	
	print "exported %d flows"%count
	if format == '0':
		outfile.close()
			
if __name__ == '__main__': main()